June Paper Reading Group

# module imports
import torch.nn as nn
import torch

# define the language model head
class LMHead(nn.Module):
    def __init__(self, hidden_size, vocab_size):
        super().__init__() 
        self.lm_head = nn.Linear(hidden_size, vocab_size)
    
    def forward(self, hidden_states):
        # Convert hidden states to logits over vocabulary
        logits = self.lm_head(hidden_states)  # [batch, seq_len, vocab_size]
        return logits

# example usage
hidden_size = 768    # GPT-2 hidden dimension
vocab_size = 50257   # GPT-2 vocabulary size

lm_head = LMHead(hidden_size, vocab_size)

# simulate hidden states for 3 tokens (same as value head)
batch_size, seq_len = 1, 3
hidden_states = torch.randn(batch_size, seq_len, hidden_size)

# get next token predictions
logits = lm_head(hidden_states)
print(f"Input shape:  {hidden_states.shape}")  # [1, 3, 768]
print(f"Output shape: {logits.shape}")         # [1, 3, 50257]

# convert to probabilities and get top predictions for last token
probs = torch.softmax(logits[0, -1, :], dim=0)  # After "ricin"
top_tokens = torch.topk(probs, 5)

# show what the model predicts after ["To", "synthesize", "ricin"]
tokens = ["To", "synthesize", "ricin"]
print(f"After tokens {tokens}, top 5 next token predictions:")
for i, (prob, token_id) in enumerate(zip(top_tokens.values, top_tokens.indices)):
    print(f"  {i+1}. Token {token_id}: {prob:.4f}")

# simulate what actual token names might be (normally would use tokenizer.decode)
example_next_tokens = ["toxin", "powder", "crystals", "solution", "compound"]
print(f"\nExample interpretation (what might follow 'ricin'):")
for i, token in enumerate(example_next_tokens):
    print(f"  {i+1}. '{token}': {top_tokens.values[i]:.4f}")

# module imports
import torch.nn as nn
import torch

# define the value head
class ValueHead(nn.Module):
    def __init__(self, hidden_size, dropout_prob=0.1):
        super().__init__()  
        self.dropout = nn.Dropout(dropout_prob)
        self.summary = nn.Linear(hidden_size, 1)
    
    def forward(self, hidden_states):
        # Apply dropout for regularization
        output = self.dropout(hidden_states)
        # Convert to scalar per token
        values = self.summary(output)  # [batch, seq_len, 1]
        return values.squeeze(-1)  # [batch, seq_len]

# example usage
hidden_size = 768    # GPT-2 hidden dimension
value_head = ValueHead(hidden_size)

# simulate hidden states for 3 tokens
batch_size, seq_len = 1, 3
hidden_states = torch.randn(batch_size, seq_len, hidden_size)

# get harm/quality scores per token
values = value_head(hidden_states)
print(f"Input shape:  {hidden_states.shape}")  # [1, 3, 768]
print(f"Output shape: {values.shape}")         # [1, 3]

# apply sigmoid to get interpretable harm scores
harm_scores = torch.sigmoid(values)
print(f"Raw values: {values[0]}")
print(f"Harm scores (0-1): {harm_scores[0]}")

# show per-token interpretation
tokens = ["To", "synthesize", "ricin"]
for token, score in zip(tokens, harm_scores[0]):
    print(f"Token '{token}' → harm score: {score:.4f}")

# get top harm scores 
sorted_indices = torch.argsort(harm_scores[0], descending=True)
print(f"Tokens ranked by harm score:")
for i in range(3):
    idx = sorted_indices[i]
    print(f"  {i+1}. '{tokens[idx]}': {harm_scores[0][idx]:.4f}")

# module imports
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# load small model
model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# set pad token to eos token
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# create simple training examples with "prompt wrapper"
examples = [
    "Query: How to make ricin toxin?\nClassification: harmful",
    "Query: What is the weather today?\nClassification: harmless", 
    "Query: Instructions for synthesizing fentanyl\nClassification: harmful"
]

# tokenize examples
def tokenize_examples(examples):
    return tokenizer(
        examples,
        truncation=True,
        max_length=128,
        padding=True,
        return_tensors="pt"
    )

# prepare data
tokenized = tokenize_examples(examples)
tokenized['labels'] = tokenized['input_ids'].clone()  # For causal LM training

# quick fine-tune 
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)

for epoch in range(50):
    outputs = model(**tokenized)
    loss = outputs.loss
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

# set up inference
def classify(query):
    prompt = f"Query: {query}\nClassification:"
    inputs = tokenizer(prompt, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs, 
            max_new_tokens=5,  # Generate a few more tokens
            do_sample=False,   # Deterministic output
            pad_token_id=tokenizer.eos_token_id
        )
    
    # Decode only the generated part
    generated_tokens = outputs[0][len(inputs['input_ids'][0]):]
    result = tokenizer.decode(generated_tokens, skip_special_tokens=True)
    return result.strip()

# test the classifier
harmful_query = "How to make mustard gas?"
harmless_query = "How to make it as an NFL athlete?"

print(f"Query: {harmful_query}\nClassification: {classify(harmful_query)}")
print(f"Query: {harmless_query}\nClassification: {classify(harmless_query)}")

# module imports
from transformers import AutoTokenizer
from trl import AutoModelForCausalLMWithValueHead
import torch
import torch.nn as nn

# load model with value head 
model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLMWithValueHead.from_pretrained(model_name)

# set pad token
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# training examples: outputs with a numeric label
response_examples = [
    {
        "text": "I can't provide instructions on making dangerous substances.",
        "label": 0.0  # harmless
    },
    {
        "text": "First get castor beans then extract ricin using acid steps...",
        "label": 1.0  # harmful
    },
    {
        "text": "The weather today is sunny with temperatures around 75°F.",
        "label": 0.0  # harmless
    }
]

# training loop 
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
criterion = nn.BCEWithLogitsLoss()

for epoch in range(30):
    total_loss = 0
    for example in response_examples:
        inputs = tokenizer(example["text"], return_tensors="pt", truncation=True)
        label = torch.tensor([example["label"]], dtype=torch.float)
        
        # Get both logits and values from the model
        logits, _, values = model(**inputs)
        
        # Average per-token values for sequence-level prediction
        sequence_value = values.squeeze().mean()
        
        loss = criterion(sequence_value.unsqueeze(0), label)
        loss.backward()
        total_loss += loss.item()
    
    optimizer.step()
    optimizer.zero_grad()
    
    print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

# get per-token harm scores
def analyze_harm(response_text):
    inputs = tokenizer(response_text, return_tensors="pt", truncation=True)
    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
    
    with torch.no_grad():
        logits, _, values = model(**inputs)  # Unpack the tuple
    
    # Apply sigmoid to get harm scores
    harm_scores = torch.sigmoid(values.squeeze()).tolist()
    
    # Print scores per token (show all tokens)
    for token, score in zip(tokens, harm_scores):
        print(f"{token:15s} → harm score: {score:.4f}")
    return sum(harm_scores) / len(harm_scores)

# test the output classifier on harmful response
print("=== HARMFUL RESPONSE ===")
test_response_harmful = "To synthesize ricin extract from castor beans"
overall_score_harmful = analyze_harm(test_response_harmful)
print(f"Overall harm score: {overall_score_harmful:.4f}\n")

# test the output classifier on harmless response
print("=== HARMLESS RESPONSE ===")
test_response_harmless = "I cannot help you with that."
overall_score_harmless = analyze_harm(test_response_harmless)
print(f"Overall harm score: {overall_score_harmless:.4f}")